"""
CSPBuilder plugin

Provides grid functionality

Copyright (c) 2008 Rune M. Friborg, runef@diku.dk.
See LICENSE.txt for licensing details (GPLv2).
"""

PLUGIN_NAME = "ForkNJoin"

try:
	import cPickle as pickle
except:
	import pickle

import base64
import mig.core
import mig.csp2mig.client
from extensions import cspxml
from plugins.execute import *
import uuid
import os

class PluginForkNJoin(PluginEXECUTION):
	
	def __init__(self, pluginManager):
		PluginEXECUTION.__init__(self, pluginManager, PLUGIN_NAME)
		self.is_grid_enabled = False
		self.is_debug_enabled = False
		self.compress_data = False		
		self.mig_contexts = {}
		self.send_to_grid_childnodes = {}
		self.vgrid = None
		# if a .miguser.conf is created then we use it. else we rely on the user to have a ~/.mig folder.
		server_loaded = False
		cert_loaded = False
		key_loaded = False
		passwd_loaded = False

		self.home = os.getenv('USERPROFILE') or os.getenv('HOME')
		self.miguser_conf = self.home + '/.cspbuilder/miguser.conf'
		self.migprefs_conf = self.home + '/.cspbuilder/migprefs.conf'
		self.local_miguser_conf = '.cspbuilder/miguser.conf'
		self.local_migprefs_conf = '.cspbuilder/migprefs.conf'
		
		
		# we check for the existence and the validity of the file.
		if os.path.exists(self.miguser_conf) or os.path.exists(self.local_miguser_conf):
			if os.path.exists(self.local_miguser_conf):
				conf = self.local_miguser_conf
			else:
				conf = self.miguser_conf
			f = open(conf,'r')
			for line in f.readlines():
				parts = line.split()
				if len(parts) == 2:
					if parts[0] == 'migserver':
						server_loaded = True
					elif parts[0] == 'certfile':
						cert_loaded = True
					elif parts[0] == 'keyfile':
						key_loaded = True
					elif parts[0] == 'password':
						passwd_loaded = True
					else:
						print 'Got unknown setting in \'.miguser.conf\': ' + line
			if server_loaded and cert_loaded and key_loaded and passwd_loaded:
				mig.core.MIG_USER_CONF = conf
			else:
				print 'Warning: .miguser.conf incomplete'
		
		self.vgrid = None
		if os.path.exists(self.migprefs_conf) or os.path.exists(self.local_migprefs_conf):
			if os.path.exists(self.local_migprefs_conf):
				conf = self.local_migprefs_conf
			else:
				conf = self.migprefs_conf
			f = open(conf,'r')
			for line in f.readlines():
				parts = line.split()
				if len(parts) == 2:
					if parts[0] == 'mig_script_path':
						mig.core.SCRIPT_PATH = parts[1]
					elif parts[0] == 'vgrid':
						self.vgrid = parts[1]
					else:
						print 'Warning: Got unknown setting in \'.migprefs.conf\': ' + line		
	
	def enable_grid(self):
		self.is_grid_enabled = True
	
	def debug_grid(self):
		self.is_debug_enabled = True
		mig.core.PRINT_DEBUG = True # set the global variable in the mig module
		mig.csp2mig.client.PRINT_DEBUG = True # and the csp2mig client module
	
	def compress_grid(self):
		self.compress_data = True
		mig.csp2mig.client.COMPRESS = True # set the global variable in the client module
			
	def create_submitters(self,multiplier=1):
		subm_name = 'MiG_Submitter_' + uuid.uuid4().hex
		if multiplier > 1:
			ioi =  subm_name+'_InOrderIn'
			ioo =  subm_name+'_InOrderOut'
			proc_name = subm_name + '_InOrder'
			cin =  proc_name + '.in0'
			cout = proc_name + '.out0'
			subm_names = []
			# Create skeleton for compositional network
			pl = {'name':proc_name,'process':proc_name, 
				  'PROCESS': 
					{'CHANNELENDS': 
						[{'direction': 'in', 'type': 'unknown', 'name': 'in0', 'target': ioi + '.in0'}, 
						 {'direction': 'out', 'type': 'unknown', 'name': 'out0', 'target': ioo + '.out0'}],
					 'PROCESSES': 
						[{'CHANNELENDS': 
							[{'direction': 'in', 'type': 'unknown', 'name': 'in0', 'target': 'MiG_InOrderInFunc.in0'},
							 {'direction': 'out', 'type': 'unknown', 'name': 'out0', 'target': 'MiG_InOrderInFunc.out0'}], 
							'CODE': {'src': 'MiG_InOrderIn', 'func': 'MiG_InOrderInFunc'},
							'name': ioi},
						{'CHANNELENDS': 
							[{'direction': 'in', 'type': 'unknown', 'name': 'in0', 'target': 'MiG_InOrderOutFunc.in0'},
							 {'direction': 'out', 'type': 'unknown', 'name': 'out0', 'target': 'MiG_InOrderOutFunc.out0'}], 
							'CODE': {'src': 'MiG_InOrderOut', 'func': 'MiG_InOrderOutFunc'},
							'name': ioo}										
						],
					 'CPOINTS' : [{'name' : 'inCpoint'} , {'name' : 'outCpoint'}],		
					 'CHANNELS' : [{'name': 'inorder1' ,'outTarget' : ioi + '.out0' ,'inTarget' : 'inCpoint.' },
								   {'name': 'inorder2' ,'outTarget' : 'outCpoint.','inTarget' : ioo+'.in0' } ],
					 'name': proc_name, 'package': 'MiG'}
				 }
			# create the number of submitters needed
			for num in range(multiplier):
				cur_subm = subm_name + '_' + str(num)
				submitter = {'CHANNELENDS': 
					[{'direction': 'in', 'type': 'unknown', 'name': 'in0', 'target': 'MiG_SubmitterFunc.in0'},
					 {'direction': 'out', 'type': 'unknown', 'name': 'out0', 'target': 'MiG_SubmitterFunc.out0'}], 
					'CODE': {'src': 'MiG_Submitter', 'func': 'MiG_SubmitterFunc'},
					'name': cur_subm}
				pl['PROCESS']['PROCESSES'].append(submitter)
				pl['PROCESS']['CHANNELS'].append({ 'name': 'ioi_' + str(num) ,'outTarget' : 'inCpoint.' ,'inTarget' : cur_subm+'.in0' })
				pl['PROCESS']['CHANNELS'].append({'name': 'ioo_' + str(num), 'outTarget' : cur_subm+'.out0','inTarget' : 'outCpoint.' })
				subm_names.append(cur_subm)
		else:
			cin =  subm_name+'.in0'
			cout =	subm_name+'.out0'
			pl = {'name':subm_name,'process':subm_name, 
				  'PROCESS': 
					{'CHANNELENDS': 
						[{'direction': 'in', 'type': 'unknown', 'name': 'in0', 'target': cin}, 
						 {'direction': 'out', 'type': 'unknown', 'name': 'out0', 'target': cout}],
					 'PROCESSES': 
						[{'CHANNELENDS': 
							[{'direction': 'in', 'type': 'unknown', 'name': 'in0', 'target': 'MiG_SubmitterFunc.in0'},
							 {'direction': 'out', 'type': 'unknown', 'name': 'out0', 'target': 'MiG_SubmitterFunc.out0'}], 
							'CODE': {'src': 'MiG_Submitter', 'func': 'MiG_SubmitterFunc'},
							'name': subm_name}],
					'name': subm_name, 'package': 'MiG'}
				 }
			subm_names = [subm_name]
		return (pl, cin, cout, subm_names)
	
		
	def replace_channels(self, proc, plink, new_in, new_out, context):
		# find old channel names
		global_proc_name = plink['name']
		old_in = ''
		old_out = ''
		if plink.has_key('PROCESS'):
			old_proc = plink['PROCESS']
			if old_proc.has_key('CHANNELENDS'):
				for chan_end in old_proc['CHANNELENDS']:
					if chan_end['direction'] == 'in':
						old_in = global_proc_name + '.' + chan_end['name']
					elif chan_end['direction'] == 'out':
						old_out = global_proc_name + '.' + chan_end['name']
		if old_in != '' and old_out != '' and proc.has_key('CHANNELS'):
			for channel in proc['CHANNELS']:
				if channel['inTarget'] == old_in:
					channel['inTarget'] = new_in
				elif channel['outTarget'] == old_out:
					channel['outTarget'] = new_out
			context['in_channel'] = old_in
			context['out_channel'] = old_out
		else:
			print 'Error: Unable to replace channels ('+old_in +','+old_out+')' +' with (' + new_in + ',' + new_out +')'
	
		
	def add_grid_procs(self, p, path='main', parent_mig_context=None):
		if not self.is_grid_enabled:
			print 'Grid not enabled. Use the \'-grid\' option to enable grid jobs. Any grid job will be run on your local machine.'
			return p
		process_name = path + '.' + p['name']
		if p.has_key('PROCESSLINKS'):
			
			# first we pick up all grid processes
			process_links = []
			i = 0
			for pl in p['PROCESSLINKS']:
				cur_mig_context = parent_mig_context
				if pl.has_key('sendToGrid') and pl['sendToGrid'] == '1':
					# we know for sure(!) that this process-link has a mig_context, because it was found in load_mig_tags
					if parent_mig_context == None:
						cur_mig_context = pl['mig_context'] 
						# contexts added to the list below will have their own submitter				
						process_links.append((i,cur_mig_context))								
					else:
						# we don't create a new submitter, but instead we append the files of the sub processlink
						# a the moment we discard any further mig-settings made by children of the parent mig process
						# this is probably an OK decision, since the programmer should be aware of the CPU-time demands
						# of any subnetwork.
						parent_mig_context['input_files'].extend(pl['mig_context']['input_files'])
						parent_mig_context['output_files'].extend(pl['mig_context']['output_files'])					
					
				# traverse sub procs of this process link and provide the current context.
				# the sub procs' contexts will be appended to the supplied context
				self.add_grid_procs(pl['PROCESS'], process_name,cur_mig_context)
				i += 1
			
			# create any submitters needed
			for index, context in process_links:
				pl = p['PROCESSLINKS'][index]

				if context.has_key('inOrderJobs'):
					num_subms = context['inOrderJobs']
				else:
					num_subms = 1
					
				# create the needed submitter(s) and add to the mig_context
				(submitter, cin_name, cout_name, subm_names) = self.create_submitters(num_subms)				 
				context['process'] = pl['process']
				context['process_name'] = pl['name']
				context['process_location'] = pl['file']						
				# replace the channels
				self.replace_channels(p, pl, cin_name, cout_name, context)
				# replace processlink
				p['PROCESSLINKS'][index] = submitter
				# store changes. they will be used in edit_param_list
				for subm_name in subm_names:
					self.mig_contexts[subm_name] = context
		if p.has_key('PROCESSES'):
			for proc in p['PROCESSES']:
				self.add_grid_procs(proc, process_name, parent_mig_context)
		return p
	
	def edit_param_list(self, process_obj, L, D):	
		parts = process_obj.name.split('.')
		short_name = parts[len(parts)-1]
		if self.mig_contexts.has_key(short_name):
			# add the changes to the parameter list
			context = self.mig_contexts[short_name]
			context['vgrid'] = self.vgrid # VGRID is not configurable on a per process basis
			L.append(context)
		return (L,D)
	
		
	def load_mig_tags(self, node, att):
		# Does processlink contain any process?
		if (att.has_key('PROCESS')):
			# check if any processlink has to be sent to grid and if so save its childnodes
			# into the attribute list.
			node_attrs = cspxml.getAttrList(node)
			if node_attrs.has_key('sendToGrid') and node_attrs['sendToGrid'] == '1':
				mig_context = {} # create a MiG-context. This hashtable will be passed around internaly
				att['mig_context'] = mig_context
				mig_context['send_to_grid_child_nodes'] = node.childNodes
				mig_context['input_files'] = []
				mig_context['output_files'] = []
				# Load files into the attributes.
				for child in node.childNodes:
					if (child.nodeName == 'inputFile'):
						attrs = cspxml.getAttrList(child)
						mig_context['input_files'].append((attrs['local_path'], attrs['remote_path']))
					elif child.nodeName == 'outputFile':
						attrs = cspxml.getAttrList(child)
						mig_context['output_files'].append((attrs['local_path'], attrs['remote_path']))
					elif child.nodeName == 'migJobConfig':
						attrs = cspxml.getAttrList(child)
						if attrs.has_key('memory'):
							mig_context['memory'] = int(attrs['memory'])
						if attrs.has_key('cpuTime'):
							mig_context['cpu_time'] = int(attrs['cpuTime'])
						if attrs.has_key('nodeCount'):
							mig_context['node_count'] = int(attrs['nodeCount'])
						if attrs.has_key('cpuCount'):
							mig_context['cpu_count'] = int(attrs['cpuCount'])
						if attrs.has_key('disk'):
							mig_context['disk'] = int(attrs['disk'])
						if attrs.has_key('architecture'):
							mig_context['arch'] = attrs['architecture']
						if attrs.has_key('inOrderJobs'):
							mig_context['inOrderJobs'] = int(attrs['inOrderJobs'])
														
		return att	
	


def load(pluginManager):
	p = PluginForkNJoin(pluginManager)
	p.addCommandLineArg('-grid', '--enable-grid', p.enable_grid, '\n  -grid, --enable-grid		Enable grid execution\n')
	p.addCommandLineArg('-dg', '--debug-grid', p.debug_grid, '  -dg, --debug-grid		Print all grid settings\n')
	p.addCommandLineArg('-cg', '--compress-grid', p.compress_grid, '  -cg, --compress-grid		Compress data sent through grid process channels\n')	
	p.addCSPTreeProcessing(p.add_grid_procs)
	p.addToProcessCreation(p.edit_param_list)
	p.addCSPFileProcessing(node_PROCESS_LINK,p.load_mig_tags)
	pluginManager.addPlugin(p)
